#ifndef RL_PLANNING_ENGINE_H
#define RL_PLANNING_ENGINE_H

#include "algohub_common_type.h"
#include "RLPlatformData.h"
#include "NDimsArray.h"
#include <chrono>
#include <random>
#include <vector>
#include <memory> // For std::unique_ptr
#include "RLRewardEngine.h"

namespace algo
{
    class RLPlanningEngine
    {
    public:
        RLPlanningEngine(const RLPlatformData& platformData, const std::map<std::string, darr::DimsArray<float>>& topographyData, const algobase::RLConfig& config);
        RLPlanningEngine(const RLPlatformData& platformData, const darr::DimsArray<float>& depthData, const algobase::RLConfig& config); // Overload for testing
        ~RLPlanningEngine();

        algobase::Path findPath(const algobase::Coord& start, const algobase::Coord& end, const std::string& output_dir = ".");
        double getTrainingDuration() const;

    public:
        // Enum for reward vector indices
        enum RewardType {
            REWARD_SAFETY = 0,
            REWARD_GOAL_APPROACH,
            REWARD_STEP_COST,
            REWARD_HEURISTIC,
            REWARD_VECTOR_SIZE
        };
        
        // State and action space
        enum Action { North, South, East, West, Up, Down, Invalid = -1 };
        static const int NUM_ACTIONS = 6;

        // Public helper methods needed by RLRewardEngine
        algobase::CoordZ gridToWorld(int ix, int iy, int iz) const;
        double getDepth(double lon, double lat) const;

    private:
        // Performance metrics
        std::chrono::duration<double> mTrainingDuration;

        // Environment parameters
        int mGridSizeX, mGridSizeY, mGridSizeZ;
        double mResolutionX, mResolutionY, mResolutionZ;
        algobase::Coord mOrigin;

        // RL parameters
        algobase::RLConfig mConfig;
        double mCurrentEpsilon; // Current exploration rate, which will decay
        std::vector<std::vector<std::vector<std::array<std::vector<double>, NUM_ACTIONS>>>> mQTable;
        std::unique_ptr<RLRewardEngine> mRewardEngine;

        const RLPlatformData& mPlatformData;
        const darr::DimsArray<float> mDepthData;
        std::mt19937 mRandGen; // Random number generator

        // Logging and output
        std::string mOutputDir;

        // Private methods
        void discretizeEnvironment(const algobase::Coord& start, const algobase::Coord& end);
        void initializeQTable();
        Action chooseAction(int ix, int iy, int iz);
        void updateQValue(int ix, int iy, int iz, Action action, const std::vector<double>& reward, int next_ix, int next_iy, int next_iz);
        std::tuple<int, int, int> worldToGrid(const algobase::CoordZ& pos) const;
        void saveIntermediatePath(int episode, const algobase::Coord& start, const algobase::Coord& end) const;
        void saveQTable(int episode) const;
        double findSafeDepth(const algobase::Coord& pos) const;
    };
}

#endif // !RL_PLANNING_ENGINE_H